import torch
from easycore.common.registry import Registry

SCHEDULER_REGISTRY = Registry('scheduler')

def build_scheduler(cfg, optimizer):
    return SCHEDULER_REGISTRY.get(cfg.SCHEDULER.NAME)(cfg, optimizer)
    
@SCHEDULER_REGISTRY.register()
def build_ReduceLROnPlateau_scheduler(cfg, optimizer):
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=cfg.SCHEDULER.PATIENCE, verbose=True)
    return scheduler

@SCHEDULER_REGISTRY.register()
def build_cosine_annealing_scheduler(cfg, optimizer):
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.SCHEDULER.T_MAX, eta_min=cfg.SCHEDULER.ETA_MIN)
    return scheduler
